#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
高级使用示例
演示如何使用 FrankaPosServo 进行高级的机器人控制
"""

import time
import numpy as np
from frankaservo import FrankaPosServo

class RobotTrajectoryController:
    """机器人轨迹控制器"""
    
    def __init__(self, robot_ip: str = "172.16.0.2"):
        self.robot = FrankaPosServo(ip=robot_ip, dynamics_factor=0.3, goal_tolerance=0.0005)
        self.home_position = self.robot.robot_position.copy()
        
    def execute_trajectory(self, waypoints, speeds=None, wait_time=0.5):
        """
        执行轨迹
        
        Args:
            waypoints: 路径点列表 (相对于当前位置)
            speeds: 对应的速度列表
            wait_time: 每个点的等待时间
        """
        if speeds is None:
            speeds = [0.01] * len(waypoints)
        
        print(f"开始执行轨迹，共 {len(waypoints)} 个路径点")
        
        for i, (waypoint, speed) in enumerate(zip(waypoints, speeds)):
            print(f"移动到路径点 {i+1}: {waypoint}, 速度: {speed}")
            
            cmd = self.robot.gen_rela_move_cmd(waypoint, max_speed=speed)
            self.robot.blocking_move(cmd)
            
            time.sleep(wait_time)
            
            # 检查当前位置
            current_pos = self.robot.robot_position
            print(f"当前位置: {current_pos}")
        
        print("轨迹执行完成")
    
    def circle_trajectory(self, center_offset, radius, num_points=8):
        """
        生成圆形轨迹
        
        Args:
            center_offset: 圆心相对于当前位置的偏移
            radius: 圆的半径
            num_points: 圆上的点数
        """
        angles = np.linspace(0, 2*np.pi, num_points, endpoint=False)
        waypoints = []
        
        for angle in angles:
            x = center_offset[0] + radius * np.cos(angle)
            y = center_offset[1] + radius * np.sin(angle)
            z = center_offset[2]
            waypoints.append([x, y, z])
        
        # 回到起点
        waypoints.append(center_offset)
        
        return waypoints
    
    def spiral_trajectory(self, center_offset, max_radius, height_change, num_turns=3):
        """
        生成螺旋轨迹
        
        Args:
            center_offset: 螺旋中心相对于当前位置的偏移
            max_radius: 最大半径
            height_change: 高度变化
            num_turns: 螺旋圈数
        """
        num_points = num_turns * 12  # 每圈12个点
        angles = np.linspace(0, 2*np.pi*num_turns, num_points)
        
        waypoints = []
        for i, angle in enumerate(angles):
            progress = i / (num_points - 1)
            radius = max_radius * progress
            
            x = center_offset[0] + radius * np.cos(angle)
            y = center_offset[1] + radius * np.sin(angle)
            z = center_offset[2] + height_change * progress
            
            waypoints.append([x, y, z])
        
        return waypoints
    
    def return_home(self):
        """返回初始位置"""
        print("返回初始位置...")
        cmd = self.robot.gen_abs_move_cmd(self.home_position, max_speed=0.02)
        self.robot.blocking_move(cmd)
        print("已返回初始位置")

def main():
    """高级使用示例主函数"""
    
    print("初始化高级机器人控制器...")
    controller = RobotTrajectoryController()
    
    try:
        # 示例 1: 圆形轨迹
        print("\n=== 示例 1: 圆形轨迹 ===")
        circle_waypoints = controller.circle_trajectory(
            center_offset=[0.03, 0.03, 0.02],
            radius=0.02,
            num_points=8
        )
        controller.execute_trajectory(circle_waypoints, wait_time=0.3)
        
        # 返回起始位置
        controller.return_home()
        time.sleep(1)
        
        # 示例 2: 螺旋轨迹
        print("\n=== 示例 2: 螺旋轨迹 ===")
        spiral_waypoints = controller.spiral_trajectory(
            center_offset=[0, 0, 0],
            max_radius=0.03,
            height_change=0.05,
            num_turns=2
        )
        speeds = [0.005 + 0.01 * (i / len(spiral_waypoints)) for i in range(len(spiral_waypoints))]
        controller.execute_trajectory(spiral_waypoints, speeds=speeds, wait_time=0.1)
        
        # 返回起始位置
        controller.return_home()
        time.sleep(1)
        
        # 示例 3: 复杂轨迹 - 写字母 "H"
        print("\n=== 示例 3: 写字母 'H' ===")
        h_waypoints = [
            [0, 0, 0.03],        # 抬起
            [-0.02, 0, 0.03],    # 移动到起点
            [-0.02, 0, 0],       # 下降
            [-0.02, 0.04, 0],    # 画左边的竖线
            [-0.02, 0.04, 0.03], # 抬起
            [-0.02, 0.02, 0.03], # 移动到中间
            [-0.02, 0.02, 0],    # 下降
            [0.02, 0.02, 0],     # 画横线
            [0.02, 0.02, 0.03],  # 抬起
            [0.02, 0, 0.03],     # 移动到右边底部
            [0.02, 0, 0],        # 下降
            [0.02, 0.04, 0],     # 画右边的竖线
            [0.02, 0.04, 0.03],  # 抬起
            [0, 0, 0.03],        # 回到中心上方
            [0, 0, 0],           # 下降到起始高度
        ]
        
        h_speeds = [0.008] * len(h_waypoints)
        controller.execute_trajectory(h_waypoints, speeds=h_speeds, wait_time=0.2)
        
        # 最终返回初始位置
        controller.return_home()
        
    except Exception as e:
        print(f"执行过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
    
    finally:
        # 确保机器人停止
        controller.robot.stop_robot()
        print("\n高级示例完成")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n用户中断程序")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
        import traceback
        traceback.print_exc()
